# 内存管理工具

import gc
import psutil
import logging
from typing import Optional
from contextlib import contextmanager

logger = logging.getLogger(__name__)

class MemoryManager:
    """
    内存管理工具类，用于监控和优化内存使用
    """
    
    def __init__(self, memory_threshold: float = 75.0):
        """
        初始化内存管理器
        
        Args:
            memory_threshold: 内存使用率阈值(%)，超过此值时触发清理
        """
        self.memory_threshold = memory_threshold
        self.process = psutil.Process()
    
    def get_memory_usage(self) -> float:
        """
        获取当前内存使用率
        
        Returns:
            内存使用率百分比
        """
        try:
            memory_info = psutil.virtual_memory()
            return memory_info.percent
        except Exception as e:
            logger.warning(f"获取内存使用率失败: {e}")
            return 0.0
    
    def get_process_memory(self) -> float:
        """
        获取当前进程内存使用量(MB)
        
        Returns:
            进程内存使用量(MB)
        """
        try:
            memory_info = self.process.memory_info()
            return memory_info.rss / 1024 / 1024  # 转换为MB
        except Exception as e:
            logger.warning(f"获取进程内存使用量失败: {e}")
            return 0.0
    
    def force_garbage_collection(self) -> int:
        """
        强制执行垃圾回收
        
        Returns:
            回收的对象数量
        """
        try:
            # 执行三代垃圾回收
            collected = 0
            for generation in range(3):
                collected += gc.collect(generation)
            
            logger.debug(f"垃圾回收完成，回收对象数: {collected}")
            return collected
        except Exception as e:
            logger.warning(f"垃圾回收失败: {e}")
            return 0
    
    def clear_gpu_cache(self):
        """
        清理GPU缓存（如果可用）
        """
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
                logger.debug("GPU缓存已清理")
        except ImportError:
            pass  # PyTorch未安装
        except Exception as e:
            logger.warning(f"清理GPU缓存失败: {e}")
    
    def check_and_cleanup(self) -> bool:
        """
        检查内存使用率，如果超过阈值则执行清理
        
        Returns:
            是否执行了清理操作
        """
        memory_usage = self.get_memory_usage()
        
        if memory_usage > self.memory_threshold:
            logger.warning(f"内存使用率过高: {memory_usage:.1f}%，开始清理")
            
            # 执行垃圾回收
            collected = self.force_garbage_collection()
            
            # 清理GPU缓存
            self.clear_gpu_cache()
            
            # 再次检查内存使用率
            new_memory_usage = self.get_memory_usage()
            logger.info(f"内存清理完成: {memory_usage:.1f}% -> {new_memory_usage:.1f}%")
            
            return True
        
        return False
    
    @contextmanager
    def memory_monitor(self, operation_name: str = "操作"):
        """
        内存监控上下文管理器
        
        Args:
            operation_name: 操作名称，用于日志记录
        """
        start_memory = self.get_process_memory()
        start_system_memory = self.get_memory_usage()
        
        logger.debug(f"{operation_name}开始 - 进程内存: {start_memory:.1f}MB, 系统内存: {start_system_memory:.1f}%")
        
        try:
            yield self
        finally:
            end_memory = self.get_process_memory()
            end_system_memory = self.get_memory_usage()
            
            memory_diff = end_memory - start_memory
            system_diff = end_system_memory - start_system_memory
            
            logger.debug(
                f"{operation_name}完成 - 进程内存变化: {memory_diff:+.1f}MB, "
                f"系统内存变化: {system_diff:+.1f}%"
            )
            
            # 如果内存使用增加过多，执行清理
            if memory_diff > 100 or end_system_memory > self.memory_threshold:
                self.check_and_cleanup()
    
    def get_memory_stats(self) -> dict:
        """
        获取详细的内存统计信息
        
        Returns:
            内存统计信息字典
        """
        try:
            # 系统内存信息
            system_memory = psutil.virtual_memory()
            
            # 进程内存信息
            process_memory = self.process.memory_info()
            
            # 垃圾回收统计
            gc_stats = gc.get_stats()
            
            return {
                "system": {
                    "total_mb": system_memory.total / 1024 / 1024,
                    "available_mb": system_memory.available / 1024 / 1024,
                    "used_mb": system_memory.used / 1024 / 1024,
                    "percent": system_memory.percent
                },
                "process": {
                    "rss_mb": process_memory.rss / 1024 / 1024,
                    "vms_mb": process_memory.vms / 1024 / 1024,
                    "percent": self.process.memory_percent()
                },
                "gc": {
                    "collections": [stat["collections"] for stat in gc_stats],
                    "collected": [stat["collected"] for stat in gc_stats],
                    "uncollectable": [stat["uncollectable"] for stat in gc_stats]
                }
            }
        except Exception as e:
            logger.error(f"获取内存统计信息失败: {e}")
            return {}

# 创建全局内存管理器实例
memory_manager = MemoryManager()